package top.yudoge.dogeeditor.app.web

import dev.langchain4j.data.segment.TextSegment
import dev.langchain4j.store.embedding.EmbeddingStore
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.web.bind.annotation.PostMapping
import org.springframework.web.bind.annotation.RequestBody
import org.springframework.web.bind.annotation.RequestMapping
import org.springframework.web.bind.annotation.RestController
import top.yudoge.dogeeditor.app.query.CompletionRequest
import top.yudoge.dogeeditor.app.vo.CompletionVO
import top.yudoge.dogeeditor.commons.utils.toJsonStr
import top.yudoge.dogeeditor.commons.web.Response
import top.yudoge.dogeeditor.db.api.DataSourceMetaDataApi
import top.yudoge.dogeeditor.llm.embedding.EmbeddingStoreHolder
import top.yudoge.dogeeditor.llm.config.models.ModelRegistry
import top.yudoge.dogeeditor.llm.objects.mro.TableMRO
import top.yudoge.dogeeditor.llm.task.TableEmbeddingTask
import top.yudoge.dogeeditor.llm.util.fetchMROByTableName
import top.yudoge.dogeeditor.llm.util.getRelatedTableNames

@RestController
@RequestMapping("/code-completion")
class CodeCompletionController {

    @Autowired
    private lateinit var tableEmbeddingTask: TableEmbeddingTask

    @Autowired
    private lateinit var metaDataApi: DataSourceMetaDataApi

    @Autowired
    private lateinit var modelRegistry: ModelRegistry

    @Autowired
    private lateinit var embeddingStoreHolder: EmbeddingStoreHolder

    @PostMapping
    fun complete(@RequestBody request: CompletionRequest): Response<CompletionVO> {
        val tables = mutableListOf<TableMRO>()
        handleTables(request, tables)
        var prompt = handlePrompt(request, tables)

        println(prompt)
        var chatModel = modelRegistry.getChatModel("qwen-plus")
        var response = chatModel.chat(prompt)
        println(response)

        var completion = response.replace("```sql", "").replace("```", "")
        return Response.success(CompletionVO(completion))
    }


    val prompt_template = """
            你是一个SQL专家，Code中是你正在编写的SQL代码其中的片段，请补全 {{CURSOR_HERE}} 处的代码。关联知识处可能是你会用到的相关知识。

            要求：
            1. 你要补的代码不应该过长，控制在3~5行
            2. 你并不需要补出完整的SQL语句
            3. 你应只返回要插入到{{CURSOR_HERE}}处的代码（包含空格等空白字符），不返回其它任何信息
            4. 如果你感觉到不知道如何补全，请放弃


            # Code
            ```sql
            #{user_code}
            ```

            # 关联知识
            #{related_knowledge}
        """


    private fun handlePrompt(
        request: CompletionRequest,
        tables: MutableList<TableMRO>
    ): String {
        var code = request.code
        var cursorPosition = request.cursorPosition
        
        // 将代码按行分割
        val lines = code.split("\n")
        
        // 计算光标所在的行号
        var currentLine = 0
        var charCount = 0
        for (i in lines.indices) {
            if (charCount + lines[i].length >= cursorPosition) {
                currentLine = i
                break
            }
            charCount += lines[i].length + 1 // +1 for newline character
        }
        
        // 提取光标所在行及前后三行
        val startLine = maxOf(0, currentLine - 3)
        val endLine = minOf(lines.size - 1, currentLine + 3)
        val targetLines = lines.subList(startLine, endLine + 1)
        
        // 计算在目标代码片段中的光标位置
        var targetCursorPosition = cursorPosition
        for (i in 0 until startLine) {
            targetCursorPosition -= lines[i].length + 1
        }
        
        // 在光标位置插入{{CURSOR_HERE}}
        val targetCodeSnippet = targetLines.joinToString("\n")
        val codeWithCursor = targetCodeSnippet.substring(0, targetCursorPosition) + 
                            "{{CURSOR_HERE}}" + 
                            targetCodeSnippet.substring(targetCursorPosition)
        
        // 准备提示词
        val relatedKnowledge = tables.toJsonStr()
        val prompt = prompt_template
            .replace("#{user_code}", codeWithCursor)
            .replace("#{related_knowledge}", relatedKnowledge)

        return prompt
    }

    private fun handleTables(
        request: CompletionRequest,
        tables: MutableList<TableMRO>
    ) {
        if (request.dataSourceIdentifier.isNotBlank() && request.databaseName.isNotBlank() && request.schemaName.isNotBlank()) {
            val embeddingStore = embeddingStoreHolder.getStore(request.dataSourceIdentifier, request.databaseName, request.schemaName)
            val queryEmbedding = tableEmbeddingTask.embedding(request.code)
            var relatedTableNames = embeddingStore.getRelatedTableNames(queryEmbedding, 3, embeddingStoreHolder)
            tables.addAll(metaDataApi.fetchMROByTableName(request.dataSourceIdentifier, request.databaseName, request.schemaName, relatedTableNames))
        }
    }
}